module List =
struct
  include List

  let wrap x = [x]

  let cons x l = x :: l

  let rec rec_remove_fast item accu = function
    | [] -> accu
    | x :: l when x == item -> rec_remove_fast item accu l
    | x :: l -> rec_remove_fast item (x :: accu) l

  let remove_fast x l = rec_remove_fast x [] l

  let sum l  = fold_left ( +  ) 0   l
  let sumf l = fold_left ( +. ) 0.0 l

  let sum_map f l  = fold_left (fun accu x -> f x +  accu) 0   l
  let sumf_map f l = fold_left (fun accu x -> f x +. accu) 0.0 l

  let rec rec_count f accu = function
    | [] -> accu
    | x :: l -> if f x then rec_count f (accu + 1) l
                else rec_count f accu l

  let count f l = rec_count f 0 l

  let rev_iter f l = fold_right (fun n accu -> f n) l ()

  let rec iter3 f l1 l2 l3 =
    match (l1, l2, l3) with
      ([], [], []) -> ()
    | (a1::l1, a2::l2, a3::l3) -> f a1 a2 a3; iter3 f l1 l2 l3
    | (_, _, _) -> invalid_arg "List.iter3"

    (* Note: not tail recursive! *)
  let rec rec_mapi f i = function
    | [] -> []
    | x :: l -> (f i x) :: rec_mapi f (i+1) l 

  let mapi f l = rec_mapi f 0 l


    (* Note: not tail recursive! *)
  let rec map3 f l1 l2 l3 =
    match (l1, l2, l3) with
      ([], [], []) -> [] 
    | (a1::l1, a2::l2, a3::l3) -> (f a1 a2 a3) :: (map3 f l1 l2 l3)
    | (_, _, _) -> invalid_arg "List.map3"

  let rec split3 = function
      [] -> ([], [], [])
    | (x,y,z)::l ->
        let (rx, ry, rz) = split3 l in (x::rx, y::ry, z::rz)

  let rec fold_left2 f l accu = match l with
      [] -> accu
    | x :: l' -> fold_left2 f l' (f x accu)

  let rec transpose l =
    try map hd l :: transpose (map tl l) with _ -> []
end

module Array =
struct
  include Array

  let sum a  = 
    let total = ref 0 in
    for i = 0 to Array.length a - 1 do
      total := !total + a.(i)
    done; !total

  let sumf a = 
    let total = ref 0.0 in
    for i = 0 to Array.length a - 1 do
      total := !total +. a.(i)
    done; !total

  let sum_map f a = 
    let total = ref 0 in
    for i = 0 to Array.length a - 1 do
      total := !total + f a.(i)
    done; !total

  let sumf_map f a = 
    let total = ref 0.0 in
    for i = 0 to Array.length a - 1 do
      total := !total +. f a.(i)
    done; !total

  let find f a = 
    let rec recfind offset =
      if offset >= Array.length a then 
        raise Not_found
      else if f a.(offset) then 
        a.(offset) 
      else 
        recfind (offset+1) in
    recfind 0

  let fold f a initial =
    let accu = ref initial in
    for i = 0 to Array.length a - 1 do
      accu := f a.(i) !accu
    done ; !accu

    (* TODO: check lengths... *)
  let map2 f a1 a2 = mapi (fun i x -> f x a2.(i)) a1

  let iter2 f a1 a2 = (* iteri (fun i x -> f x a2.(i)) a1 *)
    for i = 0 to length a1 - 1 do
      f a1.(i) a2.(i)
    done

  let rev_iter f a =
    for i = length a - 1 downto 0 do
      f a.(i)
    done

  exception Found

  let exists f a =
    try 
      for i = 0 to length a - 1 do
        if f (unsafe_get a i) then raise Found
      done ; false
    with Found -> true
 
  let for_all f a =
    try
      for i = 0 to length a - 1 do
        if not (f (unsafe_get a i)) then raise Found
      done ; true
    with Found -> false

  let count f a =
    let cnt = ref 0 in
    for i = 0 to length a - 1 do
      if f (unsafe_get a i) then incr cnt
    done ; !cnt

  let argmin a =
    let best = ref 0 in
    for i = 1 to Array.length a - 1 do
      if a.(i) < a.(!best) then
        best := i
    done ;
    !best

  let argmax a =
    let best = ref 0 in
    for i = 1 to Array.length a - 1 do
      if a.(i) > a.(!best) then
        best := i
    done ;
    !best

  let min a =
    let best = ref a.(0) in
    for i = 1 to Array.length a - 1 do
      if a.(i) < !best then
        best := a.(i)
    done ;
    !best

  let max a =
    let best = ref a.(0) in
    for i = 1 to Array.length a - 1 do
      if a.(i) > !best then
        best := a.(i)
    done ;
    !best

  let rev_inplace a =
    let l = Array.length a in
    for i = 0 to (l - 1)/2 do
      let j = l - 1 - i in 
      let tmp = a.(j) in
      a.(j) <- a.(i);
      a.(i) <- tmp
    done

  let transpose a =
    let rows = Array.length a in
    let cols = Array.length a.(0) in
    let a' = Array.create_matrix cols rows a.(0).(0) in
    for i = 0 to rows - 1 do
      for j = 0 to cols - 1 do
        a'.(j).(i) <- a.(i).(j)
      done
    done;
    a'

  let rec flattenl_helper result offset = function
  | [] -> result
  | a :: l ->
    for i = 0 to Array.length a - 1 do 
      result.(offset + i) <- a.(i)
    done;
    flattenl_helper result (offset + Array.length a) l

  let rec copy_list result offset = function
  | [] -> offset
  | x :: l -> 
    result.(offset) <- x; 
    copy_list result (offset+1) l

  let flattenl la =
    let len = sum_map List.length la in
    if len = 0 then [||]
    else
      let init = List.hd (find ((<>) []) la) in
      let a = Array.create len init in
      let offset = ref 0 in
      for i = 0 to Array.length la - 1 do
        offset := copy_list a !offset la.(i)
      done;
      a

  let filter f a =
    let y = ref [] in
    for i = Array.length a - 1 downto 0 do 
      if f a.(i) then
        y := a.(i) :: !y
    done;
    Array.of_list !y

  let partition f a =
    let y = ref [] and n = ref [] in
    for i = Array.length a - 1 downto 0 do 
      if f a.(i) then
        y := a.(i) :: !y
      else
        n := a.(i) :: !n
    done;
    (Array.of_list !y, Array.of_list !n)

  let pos_maxf a =
    let biggest = ref neg_infinity in
    let ret = ref (-1) in
    for i = 0 to Array.length a - 1 do
      if a.(i) > !biggest then begin
        ret := i;
        biggest := a.(i)
      end
    done;
    !ret

  let pos_minf a =
    let smallest = ref infinity in
    let ret = ref (-1) in
    for i = 0 to Array.length a - 1 do
      if a.(i) < !smallest then begin
        ret := i;
        smallest := a.(i)
      end
    done;
    !ret

  let rev a =
    let a' = Array.copy a in
    rev_inplace a'; a'

end

module Hashtbl =
struct
  include Hashtbl
  let keys_to_list h = Hashtbl.fold (fun x k l -> x :: l) h []
  let dfind h key default_value =
    try Hashtbl.find h key with Not_found -> default_value
end

(* Define hash set type, along with functor interface *)
module Hashset =
struct
  type 'a t = ('a, unit) Hashtbl.t
  open Hashtbl

  let create = create
  let add hs x = if not (mem hs x) then add hs x ()
  let mem hs x = mem hs x
  let iter f hs = iter (fun x () -> f x) hs
  let fold f hs = fold (fun x () l -> f x l) hs
  let to_list hs = fold (fun x l -> x :: l) hs []
  let sum_map f hs = fold (fun x accu -> f x + accu) hs 0
  let sumf_map f hs = fold (fun x accu -> f x +. accu) hs 0.0
  let filter f hs = iter (fun n -> if not (f n) then remove hs n) hs
  type 'a t' = ('a, unit) Hashtbl.t

  module type S =
    sig
      type key
      type t
      val create: int -> t
      val clear: t -> unit
      val copy: t -> t
      val add: t -> key -> unit
      val remove: t -> key -> unit
      val mem : t -> key -> bool
      val iter: (key -> unit) -> t -> unit
      val fold: (key -> 'b -> 'b) -> t -> 'b -> 'b
      val length: t -> int
      val to_list: t -> key list
      val filter: (key -> bool) -> t -> unit
    end

(* TODO -- figure out how to do this, eventually...
  module Make(H: HashedType): (S with type key = H.t) =
    struct
      module HashType = Hashtbl.Make (H)
      type key = H.t
      type hashset = key t'
      type t = hashset
      let create = HashType.create
      let clear = clear
      let copy = copy
      let add = add
      let remove = remove 
      let mem = mem 
      let iter = iter
      let fold = fold
      let length = length
      let to_list = to_list
      let filter = filter
    end
    *)
end

module String =
struct
  include String
  let rec match_beg s1 s2 i =
    if i <= 0 then 
      true
    else if String.get s1 i != String.get s2 i then
      false
    else
      match_beg s1 s2 (i-1)

  let rec match_end s1 s2 i =
    let s1l = String.length s1 in
    let s2l = String.length s2 in
    if i <= 0 then
      true
    else if s1l >= i && s2l >= i &&
        String.get s1 (s1l - i) != String.get s2 (s2l - i) then
      false
    else
      match_end s1 s2 (i-1)

  let suffix s1 s2 = match_end s1 s2 (length s2)
  let prefix s1 s2 = match_beg s1 s2 (length s2)
end

(*
 * External modules included with this library
 *)
module Heap =
struct
  include Heap
end

module Timer =
struct
  include Timer
end


(* Utility methods *)
    
let compose f g x = f(g(x))
let ($) f g x = f (g x)

let identity x = x

let f_true x = true

let rec rev_range x =
  if x < 0 then [] 
  else x :: rev_range (x-1)

let range x = List.rev (rev_range x)

(* Important math stuffs *)

(* Useful constants *)
let log_zero = neg_infinity
let log_one = 0.0

exception EmptyList

let rec dotprod a b = 
 (match a, b with
    ah :: at, bh :: bt -> (ah *. bh) +. dotprod at bt
  | [], [] -> 0.0
  | _ -> raise EmptyList)

let adotprod a b = 
  let total = ref 0.0 in
  for i = 0 to Array.length a - 1 do
    total := !total +. a.(i) *. b.(i)
  done ;
  !total

let maxf (a : float) (b : float) =
  if a > b then a else b

let minf (a : float) (b : float) =
  if a < b then a else b

let fast_exp x =
  if x > (-37.0) then exp x
  else 0.0

let logsumexp lv =
  match lv with x::[] -> x | [] -> log_zero | x::l ->
  let base = List.fold_left maxf log_zero lv in
  if base = log_zero then log_zero
  else
    let normlv = List.map (fun x -> x -. base) lv in
    let nonzero = List.filter (fun x -> x > (-37.0)) normlv in
    base +. log (List.sumf (List.map exp nonzero))

let logsumexp2 a b =
  if a > b then
    a +. log (1.0 +. exp(b -. a))
  else if a < b then
    b +. log (1.0 +. exp(a -. b))
  else
    a +. log 2.0

let alogsumexp lv =
  if Array.length lv = 1 then lv.(0) else
  let base = Array.fold_left maxf log_zero lv in
  if base = log_zero then log_zero 
  else begin
    let total = ref (fast_exp(lv.(0) -. base)) in
    for i = 1 to Array.length lv - 1 do
      total := !total +. fast_exp(lv.(i) -. base)
    done;
    base +. log(!total)
  end

let sumexp lv = exp(logsumexp lv)
let asumexp lv = exp(alogsumexp lv)

(*
let logsumexp2 a b =
  let (big, small) = if a > b then (a,b) else (b,a) in
  let diff = big -. small in
  if diff > 10.0 then big
  else big +. log (1.0 +. exp (small -. big))
 *)

let logdiffexp2 a b =
  if a -. b > 10.0 then 
    a
  else if a < b then 
    log_zero
  else 
    a +. log(1.0 -. exp(b -. a))

let output_int channel i = output_string channel (string_of_int i)

let output_float channel i = output_string channel (string_of_float i)

let output_sep channel sep a =
  for i = 0 to Array.length a - 1 do
    if i > 0 then output_string channel sep;
    output_string channel a.(i)
  done

(*
let print_objsize x = 
  let t = Objsize.objsize x in
  Printf.printf "Data: %d;  Headers: %d;  Depth: %d\n" 
    t.Objsize.data t.Objsize.headers t.Objsize.depth ;
  Printf.printf "Bytes with headers: %d\n" (Objsize.size_with_headers t) ;
  Printf.printf "Bytes without headers: %d\n" (Objsize.size_without_headers t)
  *)

(*
 * Common argument handling
 *)

(* These module-level variables shouldn't be too dangerous,
 * since there is presumably one set of arguments per program,
 * so these shouldn't get used twice.
 *)
let __verbose = ref false
let __debugging = ref false
let __logfile = ref ""
let common_arguments =
   [("-log", Arg.Set_string __logfile, " Log file") ;
    ("-v", Arg.Set __verbose, " Enable verbose output");
    ("-debug", Arg.Set __debugging, " Enable debugging output")]

(*
 * Unified logging
 *)
let log_hash = Hashtbl.create 100

let register_log (name:string) stream =
  Hashtbl.add log_hash name stream

let unregister_log name =
  Hashtbl.remove log_hash name

let log_exists name =
  Hashtbl.mem log_hash name

let log_stream name =
  Hashtbl.find log_hash name

(* Output a string to the specified log, if defined *)
let log_string name s =
  if Hashtbl.mem log_hash name then 
    output_string (Hashtbl.find log_hash name) s

(* Print to a log with printf-style formatting *)
let logf name fmt = Printf.ksprintf (log_string name) fmt

let log_normal = "normal"
let log_verbose = "verbose"
let log_debug = "debug"

let nlogf fmt = logf log_normal fmt
let vlogf fmt = logf log_verbose fmt
let dlogf fmt = logf log_debug fmt

let log_args () =
  let argstr = String.concat " " (Array.to_list Sys.argv) in
  vlogf "Arguments: %s\n" argstr

let common_log_init () =
  let debugging = !__debugging in
  let verbose = !__verbose || !__debugging in
  let logout = 
    if !__logfile = "" then stdout else open_out !__logfile in
  register_log log_normal logout;
  if verbose then register_log log_verbose logout;
  if debugging then register_log log_debug logout;
  log_args ()


(*
 * Memoization!
let memoize f get_id =
  let cache = Hashtbl.create 100 in
  let rec f' arg = (* STOP STOP -- how to make this work recursively? *)
    let id = get_id arg in
    if Hashtbl.mem cache id then
      Hashtbl.find cache id
    else begin
      let result = f arg in
      Hashtbl.add cache id result ;
      result
    end
 *)

 (*
  * Debug stuff
  *)

let debug s = print_string s; print_string "\n"; flush stdout

let string_of_ilist l = String.concat " " (List.map string_of_int l)
let string_of_flist l = String.concat " " (List.map string_of_float l)
let string_of_iarray a = string_of_ilist (Array.to_list a)
let string_of_farray a = string_of_flist (Array.to_list a)


 (*
  * Probability stuff
  *)

let normalize_inplace_raw a =
  let total = Array.sumf a in
  for i = 0 to Array.length a - 1 do
    a.(i) <- a.(i) /. total
  done

let normalize_inplace_log a =
  let logz = alogsumexp a in
  for i = 0 to Array.length a - 1 do
    a.(i) <- a.(i) -. logz
  done

let normalize_raw a = 
  let a = Array.copy a in
  normalize_inplace_raw a;
  a

let normalize_log a = 
  let a = Array.copy a in
  normalize_inplace_log a;
  a

(* Keep track of statistics conveniently *)
let stats_make () = (ref 0., ref 0., ref 0.)

let stats_wadd (sum, sumsq, n) w x =
  sum := !sum +. x *. w;
  sumsq := !sumsq +. x *. x *. w;
  n := !n +. w

let stats_add s x = stats_wadd s 1. x

let stats_n (sum, sumsq, n) = 
  !n

let stats_sum (sum, sumsq, n) = 
  !sum

let stats_mean (sum, sumsq, n) = 
  if !n = 0. then 0.
  else !sum /. !n

let stats_var (sum, sumsq, n) = 
  if !n = 0. then 0.
  else (!sumsq -. !sum *. !sum /. !n) /. !n

let stats_epsilon = 1.0e-10

let stats_stddev s = 
  let v = stats_var s in
  if v < 0.0 && v > -.stats_epsilon then
    0.0
  else
    sqrt(stats_var s)

let stats_stderr (sum, sumsq, n) =
  (stats_stddev (sum, sumsq, n)) /. (sqrt !n)


(* TODO -- is this the right place for this? *)
type libra_filetype = ACFile | BNFile | MNFile | DataFile | UnknownFile

let filetype filename =
  let filename = String.lowercase filename in
  if String.suffix filename ".ac" then ACFile
  else if String.suffix filename ".xmod" then BNFile
  else if String.suffix filename ".bif" then BNFile
  else if String.suffix filename ".dn" then BNFile
  else if String.suffix filename ".bn" then BNFile
  else if String.suffix filename ".mn" then MNFile
  else if String.suffix filename ".uai" then MNFile
  else if String.suffix filename ".simple" then MNFile
  else if String.suffix filename ".data" then DataFile
  else UnknownFile
